package digitalsignature;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

public class DigitalSignature
{
    public static String PK_FILE = "rsa_ver.bin";

    public static String SK_FILE = "rsa_sign.bin";

    /**
     * A signing algorithm that, given a message and a private key, produces a
     * signature.
     * 
     * A signature verifying algorithm that, given a message, public key and a
     * signature, either accepts or rejects the message's claim to authenticity
     * 
     * @param args
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws IOException
     * @throws InvalidKeyException
     * @throws SignatureException
     */
    public static void main( String[] args ) throws NoSuchAlgorithmException,
            InvalidKeySpecException, IOException, InvalidKeyException,
            SignatureException {
        if( args.length != 3 ) {
            System.out
                    .println( "2 params: [0, 1] - sign or ver. the message file name. the signature file name" );
            System.exit( 0 );
        }

        if( args[0].equals( "0" ) ) {
            // sign
            PublicKey pk;
            PrivateKey sk;
            String plainTextFileName = args[1];
            String signatureFileName = args[2];
            try {
                byte[] c_pkcs8 = getBytesFromFile( SK_FILE );
                sk = getPrivateKeyFromRawBytes( c_pkcs8 );
            } catch( IOException e ) {
                // generate key pair, and write file.
                KeyPairGenerator gen = KeyPairGenerator.getInstance( "RSA" );
                gen.initialize( 3072 );
                KeyPair keyPair = gen.generateKeyPair();
                pk = keyPair.getPublic();
                sk = keyPair.getPrivate();
                writePrivateKeyToFile( sk );
                writePublicKeytoFile( pk );
            }
            byte[] plainText = getBytesFromFile( plainTextFileName );
            byte[] signature = sign( plainText, sk );
            writeBytesToFile( signature, signatureFileName );
        } else if( args[0].equals( "1" ) ) {
            // verify
            PublicKey pk;
            String plainTextFileName = args[1];
            String signatureFileName = args[2];
            byte[] c_x509 = getBytesFromFile( PK_FILE );
            pk = getPublicKeyFromRawBytes( c_x509 );
            byte[] plainText = getBytesFromFile( plainTextFileName );
            byte[] signature = getBytesFromFile( signatureFileName );

            if( ver( plainText, signature, pk ) ) {
                System.out.println( "Verification Passed" );
            } else {
                System.out.println( "Verification Failed" );
            }
        } else {
            System.out.println( "wrong argument 1" );
            System.exit( 0 );
        }

    }

    private static void writeBytesToFile( byte[] bytes, String fileName )
            throws IOException {
        File file = new File( fileName );
        FileOutputStream fos = new FileOutputStream( file );
        fos.write( bytes );
        fos.close();
    }

    private static byte[] getBytesFromFile( String fileName )
            throws IOException {
        File file = new File( fileName );
        FileInputStream fis = new FileInputStream( fileName );
        byte[] data = new byte[(int) file.length()];
        fis.read( data );
        fis.close();
        return data;
    }

    /**
     * Write sk x509 encoded to file
     * 
     * @param sk
     * @throws IOException
     */
    public static void writePrivateKeyToFile( PrivateKey sk )
            throws IOException {
        System.out.print( sk.getFormat() );
        writeBytesToFile( sk.getEncoded(), SK_FILE );
    }

    /**
     * Write pk pkcs8 encoded to file
     * 
     * @param pk
     * @throws IOException
     */
    public static void writePublicKeytoFile( PublicKey pk ) throws IOException {
        System.out.print( pk.getFormat() );
        writeBytesToFile( pk.getEncoded(), PK_FILE );
    }

    /**
     * Sign plain text with private key
     * 
     * @param plainText
     * @param sk
     * @return signature
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     * @throws SignatureException
     */
    public static byte[] sign( byte[] plainText, PrivateKey sk )
            throws NoSuchAlgorithmException, InvalidKeyException,
            SignatureException {
        byte[] signature = null;
        Signature sign = Signature.getInstance( "SHA256WITHRSA" );
        sign.initSign( sk );
        sign.update( plainText );
        signature = sign.sign();
        return signature;
    }

    /**
     * Verify plain text with public key.
     * 
     * @param plainText
     * @param signature
     * @param pk
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     * @throws SignatureException
     */
    public static boolean ver( byte[] plainText, byte[] signature, PublicKey pk )
            throws NoSuchAlgorithmException, InvalidKeyException,
            SignatureException {
        Signature ver = Signature.getInstance( "SHA256WITHRSA" );
        ver.initVerify( pk );
        ver.update( plainText );
        return ver.verify( signature );
    }

    /**
     * Construct private key from pkcs8 encoded bytes.
     * 
     * @param c_pkcs8
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static PrivateKey getPrivateKeyFromRawBytes( byte[] c_pkcs8 )
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance( "RSA" );
        PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec( c_pkcs8 );
        PrivateKey privateKey = keyFactory.generatePrivate( privateKeySpec );
        return privateKey;
    }

    /**
     * Construct public key from x509 encoded bytes.
     * 
     * @param c_x509
     * @return
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static PublicKey getPublicKeyFromRawBytes( byte[] c_x509 )
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance( "RSA" );
        X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec( c_x509 );
        PublicKey publicKey = keyFactory.generatePublic( publicKeySpec );
        return publicKey;
    }
}
